import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics
import os

VER = 'v1.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)


class ConvCell(keras.Model):

    def __init__(self, filters, ksize=(3, 3), strides=(1, 1), padding='same', **kwargs):
        super().__init__(**kwargs)
        self.conv = layers.Conv2D(filters, ksize, strides, padding, use_bias=False)
        self.bn = layers.BatchNormalization()
        self.relu = layers.ReLU()

    def call(self, inputs, training=None, mask=None):
        x = self.conv(inputs, training=training)
        x = self.bn(x, training=training)
        x = self.relu(x, training=training)
        return x


class Darknet19(keras.Model):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)

        self.main_cfg = [
            32, 'm',
            64, 'm',
            128, [64, 1], 128, 'm',
            256, [128, 1], 256, 'm',
            512, [256, 1], 512, [256, 1], 512, 'm',
            1024, [512, 1], 1024, [512, 1], 1024,
            [n_cls, 1]
        ]

        self.convs = []
        for i, item in enumerate(self.main_cfg):
            print(i, item)
            if type(item) == str:
                layer = layers.MaxPool2D((2, 2), (2, 2), 'same')
            elif type(item) == int:
                layer = ConvCell(item)
            else:
                filters = item[0]
                kside = item[1]
                print('filters', filters, 'kside', kside)
                layer = ConvCell(filters, (kside, kside))
            self.convs.append(layer)

    def call(self, inputs, training=None, mask=None):
        x = inputs
        for i, layer in enumerate(self.convs):
            print(i, self.main_cfg[i], x.shape, '=>', end=' ')
            x = layer(x, training=training)
            print(x.shape)
        x = layers.GlobalAvgPool2D()(x, training=training)
        print('avg', x.shape)
        x = layers.Softmax()(x, training=training)
        print('softmax', x.shape)
        return x


model = Darknet19(2)
model.build(input_shape=(None, 224, 224, 3))
model.summary()

x = tf.zeros((4, 224, 224, 3), dtype=tf.float32)
pred = model(x)
print('pred', pred.shape)
